import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
import logging
import os
from datetime import datetime
from typing import Tuple, Dict
import tkinter as tk
import matplotlib.font_manager as fm
font_path = 'C:/Windows/Fonts/simhei.ttf'
zh_font = fm.FontProperties(fname=font_path)
plt.rcParams['font.family'] = zh_font.get_name()
plt.rcParams['axes.unicode_minus'] = False
# ===================== 全局配置 =====================
CONFIG = {
    "data_path": r"C:/智能螺栓/螺栓实时数据监测.csv",
    "batch_size": 256,
    "epochs": 3000,
    "lr": 1e-3,
    "weight_decay": 1e-2,
    "save_dir": "C:/智能螺栓/bolt_model_v7",
    "device": "cuda" if torch.cuda.is_available() else "cpu",
    "plot_interval": 50 
}
plt.switch_backend('TkAgg')
matplotlib_backend = plt.get_backend()
print(f"Using Matplotlib Backend: {matplotlib_backend}")
def setup_logger():
    logger = logging.getLogger("BoltTrainer")
    logger.setLevel(logging.INFO)    
    log_dir = CONFIG["save_dir"]
    os.makedirs(log_dir, exist_ok=True)    
    file_handler = logging.FileHandler(os.path.join(log_dir, "train.log"))
    stream_handler = logging.StreamHandler()   
    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
    file_handler.setFormatter(formatter)
    stream_handler.setFormatter(formatter)  
    logger.addHandler(file_handler)
    logger.addHandler(stream_handler)
    return logger
LOG = setup_logger()

# ===================== 数据集 =====================
class BoltDataset(Dataset):
    def __init__(self):
        self.df = pd.read_csv(CONFIG["data_path"])
        self._preprocess()
        self._init_normalizers()

    def _preprocess(self):
        q1_f = self.df["预紧力F(N)"].quantile(0.25)
        q3_f = self.df["预紧力F(N)"].quantile(0.75)
        iqr_f = q3_f - q1_f
        self.df = self.df[(self.df["预紧力F(N)"] >= q1_f - 1.5 * iqr_f) & 
                         (self.df["预紧力F(N)"] <= q3_f + 1.5 * iqr_f)]        
        q1_v = self.df["实测电压V(V)"].quantile(0.25)
        q3_v = self.df["实测电压V(V)"].quantile(0.75)
        iqr_v = q3_v - q1_v
        self.df = self.df[(self.df["实测电压V(V)"] >= q1_v - 1.5 * iqr_v) & 
                         (self.df["实测电压V(V)"] <= q3_v + 1.5 * iqr_v)]        
        self.df.reset_index(drop=True, inplace=True)
    def _init_normalizers(self):
        self.V_mean = self.df["实测电压V(V)"].mean()
        self.V_std = self.df["实测电压V(V)"].std() or 1.0
        self.F_mean = self.df["预紧力F(N)"].mean()
        self.F_std = self.df["预紧力F(N)"].std() or 1.0
    def normalize_V(self, V: np.ndarray) -> np.ndarray:
        return (V - self.V_mean) / self.V_std
    def denormalize_F(self, F_norm: np.ndarray) -> np.ndarray:
        return F_norm * self.F_std + self.F_mean
    def __len__(self):
        return len(self.df)
    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, ...]:
        row = self.df.iloc[idx]
        V_norm = self.normalize_V(np.array([row["实测电压V(V)"]]))[0]
        F_norm = (row["预紧力F(N)"] - self.F_mean) / self.F_std
        
        A = row["横截面积A(m²)"]
        E = row["弹性模量E(Pa)"]
        g31 = row["压电常数g31(V·m/N)"]
        t = row["压电片厚度t(m)"]
        V_raw = row["实测电压V(V)"]        
        return (
            torch.tensor([V_norm], dtype=torch.float32),
            torch.tensor([F_norm], dtype=torch.float32),
            torch.tensor([A], dtype=torch.float32),
            torch.tensor([E], dtype=torch.float32),
            torch.tensor([g31], dtype=torch.float32),
            torch.tensor([t], dtype=torch.float32),
            torch.tensor([V_raw], dtype=torch.float32)
        )
# ===================== PINNs环节 =====================
class PhysicsInformedNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(1, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 64),
            torch.nn.ReLU(),
            torch.nn.Linear(64, 1)
        )
        self.apply(self._init_weights)

    def _init_weights(self, m: torch.nn.Module):
        if isinstance(m, torch.nn.Linear):
            torch.nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
            torch.nn.init.zeros_(m.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(x)

    @staticmethod
    def physics_equation(F_pred_raw: torch.Tensor, A: torch.Tensor, 
                         E: torch.Tensor, g31: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        return (g31 * F_pred_raw * t) / (A * E)

# ===================== 实时绘图环节 =====================
class RealTimePlotter:
    def __init__(self):
        self.tk_root = tk.Tk()
        self.tk_root.withdraw()        
        self.fig, (self.ax1, self.ax2) = plt.subplots(2, 1, figsize=(12, 8))
        self.ax1.set_ylabel('电压 (V)')
        self.ax2.set_ylabel('预紧力 (N)')
        self.ax2.set_xlabel('训练轮次')
        self.ax1.set_title('实测电压 vs 理论电压')
        self.ax2.set_title('预测预紧力')
        self.V_real_line, = self.ax1.plot([], [], 'bo-', label='实测电压')
        self.V_pred_line, = self.ax1.plot([], [], 'r--', label='理论电压')
        self.F_pred_line, = self.ax2.plot([], [], 'g-o', label='预测预紧力')
        self.ax1.legend()
        self.ax2.legend()
        self.epoch_data = []
        self.V_real_data = []
        self.V_pred_data = []
        self.F_pred_data = []
        plt.ion() 
        plt.show(block=False)
        
    def update(self, epoch, V_real, V_pred, F_pred):
        self.epoch_data.append(epoch)
        self.V_real_data.append(V_real)
        self.V_pred_data.append(V_pred)
        self.F_pred_data.append(F_pred)
        if len(self.epoch_data) > 100:
            for lst in [self.epoch_data, self.V_real_data, self.V_pred_data, self.F_pred_data]:
                lst.pop(0)        
        # 实时更新
        self.V_real_line.set_data(self.epoch_data, self.V_real_data)
        self.V_pred_line.set_data(self.epoch_data, self.V_pred_data)
        self.F_pred_line.set_data(self.epoch_data, self.F_pred_data)
        
        self.ax1.relim()
        self.ax1.autoscale_view()
        self.ax2.relim()
        self.ax2.autoscale_view()
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
        plt.pause(0.001)  
        
    def close(self):
        plt.close(self.fig)
        self.tk_root.destroy()

# ===================== 训练引擎 =====================
class Trainer:
    def __init__(self):
        self.dataset = BoltDataset()
        self.loader = DataLoader(self.dataset, batch_size=CONFIG["batch_size"], shuffle=True)
        self.model = PhysicsInformedNN().to(CONFIG["device"])
        self.optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=CONFIG["lr"],
            weight_decay=CONFIG["weight_decay"]
        )
        self.best_loss = float("inf")
        self.plotter = RealTimePlotter() if plt.get_backend() == 'TkAgg' else None

    def train(self):
        LOG.info("开始训练，数据量：%d", len(self.dataset))
        try:
            for epoch in range(1, CONFIG["epochs"] + 1):
                total_loss = 0.0
                data_loss_sum = 0.0
                physics_loss_sum = 0.0
                
                for batch in self.loader:
                    V, F, A, E, g31, t, V_raw = [b.to(CONFIG["device"]) for b in batch]
                    
                    F_pred = self.model(V)
                    data_loss = torch.nn.functional.huber_loss(F_pred, F, delta=0.1)
                    
                    F_pred_raw = self.dataset.denormalize_F(F_pred.detach().cpu().numpy())
                    F_pred_raw = torch.tensor(F_pred_raw, device=CONFIG["device"])
                    
                    V_pred_physics = PhysicsInformedNN.physics_equation(F_pred_raw, A, E, g31, t)
                    physics_loss = torch.mean((V_raw - V_pred_physics) ** 2)
                    
                    loss = data_loss + 5 * physics_loss
                    
                    self.optimizer.zero_grad()
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                    self.optimizer.step()
                    
                    total_loss += loss.item() * len(V)
                    data_loss_sum += data_loss.item() * len(V)
                    physics_loss_sum += physics_loss.item() * len(V)
                
                avg_loss = total_loss / len(self.dataset)
                avg_data_loss = data_loss_sum / len(self.dataset)
                avg_physics_loss = physics_loss_sum / len(self.dataset)
                
                LOG.info(
                    "Epoch %4d | 总损失: %8.4e | 数据损失: %8.4e | 物理损失: %8.4e",
                    epoch, avg_loss, avg_data_loss, avg_physics_loss
                )
                
                # 实时绘图更新
                if self.plotter and epoch % CONFIG["plot_interval"] == 0:
                    with torch.no_grad():
                        batch = next(iter(self.loader))
                        V, F, A, E, g31, t, V_raw = [b.to(CONFIG["device"]) for b in batch]
                        F_pred = self.model(V)
                        F_pred_raw = self.dataset.denormalize_F(F_pred.detach().cpu().numpy())
                        V_pred_physics = PhysicsInformedNN.physics_equation(
                            torch.tensor(F_pred_raw, device=CONFIG["device"]), A, E, g31, t
                        )
                        self.plotter.update(
                            epoch, 
                            V_raw.cpu().numpy()[0], 
                            V_pred_physics.cpu().numpy()[0], 
                            F_pred_raw[0]
                        )
                
                # 保存模型
                if avg_loss < self.best_loss:
                    self.best_loss = avg_loss
                    self._save_checkpoint(epoch)
                    
        except Exception as e:
            LOG.error("训练过程中发生错误: %s", str(e))
            raise
        finally:
            if self.plotter:
                self.plotter.close() 

    def _save_checkpoint(self, epoch: int):
        state = {
            "epoch": epoch,
            "model_state_dict": self.model.state_dict(),
            "V_mean": self.dataset.V_mean,
            "V_std": self.dataset.V_std,
            "F_mean": self.dataset.F_mean,
            "F_std": self.dataset.F_std
        }
        torch.save(state, os.path.join(CONFIG["save_dir"], f"epoch_{epoch}.pth"))
        LOG.info("保存模型检查点，epoch: %d", epoch)

# ===================== 主程序（关键改进） =====================
if __name__ == "__main__":
    try:
        trainer = Trainer()
        trainer.train()
        LOG.info("训练完成，模型已保存至: %s", CONFIG["save_dir"])
    except Exception as e:
        LOG.error("训练失败: %s", str(e))
    finally:
        if 'tk' in locals() and hasattr(tk, '_default_root') and tk._default_root is not None:
            LOG.info("绘图窗口保持打开，点击关闭按钮退出")
            tk._default_root.mainloop()